package com.weizhu.swufer.core.rpc;


import com.google.common.util.concurrent.RateLimiter;
import com.weizhu.swufer.constant.SwuferConstant;
import com.weizhu.swufer.core.exception.RpcException;

import java.util.HashMap;
import java.util.Map;

public class RpcLimitManager {

    private final static Object lock = new Object();

    private final static Map<String, RateLimiter> rateLimiterMap = new HashMap<>();

    public static void acquire(RpcRequest rpcRequest){

        if (rpcRequest.getContextParameters().get("qps") == null){
            return;
        }
        String key = rpcRequest.getInterfaceName() + SwuferConstant.IP_ADDRESS_DELIMITER + rpcRequest.getMethodName();
        if(!rateLimiterMap.containsKey(key)) {
            synchronized (lock) {
                if(!rateLimiterMap.containsKey(key)) {
                    final RateLimiter rateLimiter = RateLimiter.create((Integer)rpcRequest.getContextParameters().get("qps"));
                    rateLimiterMap.put(key, rateLimiter);
                }
            }
        }
        else {
            RateLimiter rateLimiter=rateLimiterMap.get(key);
            rateLimiter.acquire();
        }
    }

    public static void tryAcquire(RpcRequest rpcRequest,int qps){
        if (qps <= 0){
            return;
        }
        String key = rpcRequest.getInterfaceName() + SwuferConstant.IP_ADDRESS_DELIMITER + rpcRequest.getMethodName();
        RateLimiter rateLimiter = null;
        if(!rateLimiterMap.containsKey(key)) {
            synchronized (lock) {
                if(!rateLimiterMap.containsKey(key)) {
                    rateLimiter = RateLimiter.create(qps);
                    rateLimiterMap.put(key, rateLimiter);
                }
            }
        } else {
            rateLimiter = rateLimiterMap.get(key);
        }
        if (rateLimiter == null){
            return;
        }
//        rateLimiter.acquire();
        if (!rateLimiter.tryAcquire()){
            throw new RpcException(RpcException.SERVER_REFUSE,
                    "Too many visitors to the server. Please try again later");
        }
    }
}
